"""
数据库连接管理
"""
from sqlalchemy import create_engine, MetaData
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import sessionmaker
from config.settings import settings
from utils.logger import audit_logger

# 创建数据库引擎
engine = create_engine(
    settings.database_url,
    pool_pre_ping=True,
    pool_recycle=300,
    echo=False  # 设置为True可以看到SQL语句
)

# 创建会话工厂
SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)

# 创建基础模型类
Base = declarative_base()

# 元数据
metadata = MetaData()


def get_db():
    """获取数据库会话"""
    db = SessionLocal()
    try:
        yield db
    except Exception as e:
        audit_logger.error(f"数据库操作异常: {e}")
        db.rollback()
        raise
    finally:
        db.close()


def init_database():
    """初始化数据库"""
    try:
        # 创建所有表
        Base.metadata.create_all(bind=engine)
        audit_logger.info("数据库初始化成功")
    except Exception as e:
        audit_logger.error(f"数据库初始化失败: {e}")
        raise


def test_connection():
    """测试数据库连接"""
    try:
        from sqlalchemy import text
        with engine.connect() as conn:
            result = conn.execute(text("SELECT 1"))
            audit_logger.info("数据库连接测试成功")
            return True
    except Exception as e:
        audit_logger.error(f"数据库连接测试失败: {e}")
        return False
